// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: leon

#pragma once
#include <map>
#include <memory>
#include <utility>
#include <vector>
#include <mutex>

#include "coin2/exchange/base/executor/topic_distributor.h"
#include "coin2/exchange/base/feed/latency_extractor.h"
#include "coin2/exchange/base/market/market_exchange_api.h"
#include "coin2/exchange/factory/feed/latency_extractor.h"
#include "coin2/exchange/base/feed/ws/feed_subscription_spec.h"
#include "coin2/exchange/feed.pb.h"

namespace coin2::exchange::base::feed {
namespace impl {

using coin2::exchange::feed::WebsocketConfig;

class LatencyManager {
 public:
  enum LatencyOption {
    NONE,
    SLOWEST,
    TIMEOUT
  };

  LatencyManager(
      const MarketExchangeApi& mea,
      const WebsocketConfig& config,
      boost::asio::io_context* ioc)
      : mea_(mea),
        num_workers_(config.num_workers()),
        strand_(*ioc) {
    if (config.reconnect_slowest()) {
      mode_ = LatencyOption::SLOWEST;
      CHECK_GT(num_workers_, 1);
    } else if (config.monitor_connection()) {
      mode_ = LatencyOption::TIMEOUT;
      timeout_ = config.monitor_timeout();
    } else {
      mode_ = LatencyOption::NONE;
    }
  }

  virtual ~LatencyManager() = default;

  std::shared_ptr<LatencyExtractor> AddLatencyExtractor(const WsConnectionConfig* config) {
    std::scoped_lock _(mutex_);
    int connection_id = config->connection_id;
    int worker_id = config->worker_id;
    if (latency_extractors_.count(connection_id) == 0) {
      latency_extractors_[connection_id] = std::map<int, std::shared_ptr<LatencyExtractor>>();
    }

    if (mode_ == LatencyOption::SLOWEST) {
      latency_extractors_[connection_id][worker_id] =
          coin2::exchange::factory::feed::NewLatencyExtractor(mea_);
    } else if (mode_ == LatencyOption::TIMEOUT) {
      latency_extractors_[connection_id][worker_id] =
          std::make_shared<coin2::exchange::base::feed::TimeoutLatencyExtractor>();
    } else {
      CHECK(false) << "Not initialized.";
    }
    
    latency_extractors_[connection_id][worker_id]->SetIoc(&strand_);
    std::pair<int, int> k(connection_id, worker_id);
    latency_connections_[k] = config;
    return latency_extractors_[connection_id][worker_id];
  }

  void RemoveLatencyExtractor(const WsConnectionConfig* config) {
    std::scoped_lock _(mutex_);
    int connection_id = config->connection_id;
    int worker_id = config->worker_id;
    latency_extractors_[connection_id].erase(worker_id);
    std::pair<int, int> k(connection_id, worker_id);
    latency_connections_.erase(k);
  }

  std::vector<const WsConnectionConfig*> RaceConnectionLatency() {
    std::scoped_lock _(mutex_);
    std::vector<const WsConnectionConfig*> ret;
    if (mode_ == LatencyOption::SLOWEST) {
      ret = RaceSlowestLatency();
    } else if (mode_ == LatencyOption::TIMEOUT) {
      ret = RaceTimeoutLatency();
    } else {
      CHECK(false) << "Not initialized.";
    }

    return ret;
  }

  bool Enabled() const {
    return mode_ != LatencyOption::NONE;
  }

  std::map<std::pair<int, int>, std::shared_ptr<TopicRecordPacket>>
  CollectConnectionInfo() const {
    std::scoped_lock _(mutex_);
    std::map<std::pair<int, int>, std::shared_ptr<TopicRecordPacket>> packets;
    for (auto const& item : latency_extractors_) {
      auto connection_id = item.first;
      auto const& map = item.second;
      for (const auto& [worker, latency_extractor] : map) {
        std::shared_ptr<TopicRecordPacket> packet = 
            latency_extractor->GetLastPacket();
        std::pair<int, int> k(connection_id, worker);
        packets[k] = packet;
      }
    }

    return packets;
  }

 private:
  std::vector<const WsConnectionConfig*> RaceSlowestLatency() {
    std::vector<const WsConnectionConfig*> ret;
    for (auto const& item : latency_extractors_) {
      auto connection_id = item.first;
      auto worker_id = 0;
      auto const& map = item.second;
      [&] {
        if (map.size() < num_workers_) {
          return;
        }
        int32_t slowest_worker = -1;
        double biggest_latency = 0;
        for (const auto& [worker, latency_extractor] : map) {
          double latency = 0;
          if (!latency_extractor->Latency(&latency)) {
            return;
          }
          if (latency > biggest_latency) {
            biggest_latency = latency;
            slowest_worker = worker;
          }
        }

        if (slowest_worker >= 0) {
          DLOG(INFO) << mea_.String() << " connection_id=" << connection_id << ", worker="
                    << slowest_worker << ", biggest_latency=" << biggest_latency;
          worker_id = slowest_worker;
          std::pair<int, int> k(connection_id, worker_id);
          ret.emplace_back(latency_connections_[k]);
        } else {
          LOG(INFO) << mea_.String() << " no latency data";
        }
      }();
    }
    return ret;
  }

  std::vector<const WsConnectionConfig*> RaceTimeoutLatency() {
    std::vector<const WsConnectionConfig*> ret;
    for (auto const& item : latency_extractors_) {
      auto connection_id = item.first;
      auto const& map = item.second;
      for (const auto& [worker, latency_extractor] : map) {
        double latency = 0;
        if (!latency_extractor->Latency(&latency)) {
          continue;
        }
        if (latency > (double)timeout_) {
          std::pair<int, int> k(connection_id, worker);
          const WsConnectionConfig* config = latency_connections_[k];
          ret.emplace_back(config);

          DLOG(INFO) << config->topic << "_" << config->url.name << "_" << config->group
              << config->worker_id << "_" << config->connection_id
              <<  " force close, no_feed_duration=" << latency;
        }
      }
    }
    return ret;
  }

  // health monitor reand, subsriber write
  mutable std::mutex mutex_;
  std::map<int, std::map<int, std::shared_ptr<LatencyExtractor>>> latency_extractors_;
  std::map<std::pair<int, int>, const WsConnectionConfig*> latency_connections_;
  MarketExchangeApi mea_;
  int32_t num_workers_;
  boost::asio::io_context::strand strand_;  // packet come in with ordered
  LatencyOption mode_;
  int64_t timeout_;  // nano seconds
};

}  // namespace impl

using impl::LatencyManager;
}  // namespace coin2::exchange::base::feed
